import tf from '../tf'
const conv = (numGrowCh: number) =>
  tf.layers.conv2d({
    filters: numGrowCh,
    kernelSize: 3,
    strides: 1,
    // dataFormat: 'channelsFirst',
    padding: 'same',
    kernelInitializer: 'glorotUniform',
    // weights,
    useBias: true,
  })

export class ResidualDenseBlock extends tf.layers.Layer {
  // implements AddWeightAgent
  numFeat: number
  numGrowCh: number
  private conv1: tf.layers.Layer
  private conv2: tf.layers.Layer
  private conv3: tf.layers.Layer
  private conv4: tf.layers.Layer
  private conv5: tf.layers.Layer
  private lrelu: (x: any) => any
  constructor(numFeat = 64, numGrowCh = 32) {
    super({})
    this.numFeat = numFeat
    this.numGrowCh = numGrowCh
    // Define convolutional layers
    this.conv1 = conv(this.numGrowCh)
    this.conv2 = conv(this.numGrowCh)
    this.conv3 = conv(this.numGrowCh)
    this.conv4 = conv(this.numGrowCh)
    this.conv5 = conv(this.numFeat)
    // LeakyReLU activation function
    const leakyRelu = tf.layers.leakyReLU({ alpha: 0.2 })
    this.lrelu = (x) => leakyRelu.apply(x)
  }
  override build(_: tf.Shape): void {
    let trainableWeights: tf.LayerVariable[] = []
    this.conv1.build([null, 64, 64, this.numFeat])
    this.conv2.build([null, 64, 64, this.numFeat + this.numGrowCh])
    this.conv3.build([null, 64, 64, this.numFeat + 2 * this.numGrowCh])
    this.conv4.build([null, 64, 64, this.numFeat + 3 * this.numGrowCh])
    this.conv5.build([null, 64, 64, this.numFeat + 4 * this.numGrowCh])
    trainableWeights.push(...this.conv1.trainableWeights)
    trainableWeights.push(...this.conv2.trainableWeights)
    trainableWeights.push(...this.conv3.trainableWeights)
    trainableWeights.push(...this.conv4.trainableWeights)
    trainableWeights.push(...this.conv5.trainableWeights)
    this._trainableWeights = trainableWeights
    this.built = true
  }
  getWeights(trainableOnly?: boolean): tf.Tensor[] {
    return [
      ...this.conv1.getWeights(trainableOnly),
      ...this.conv2.getWeights(trainableOnly),
      ...this.conv3.getWeights(trainableOnly),
      ...this.conv4.getWeights(trainableOnly),
      ...this.conv5.getWeights(trainableOnly),
    ]
  }
  setWeights(weights: tf.Tensor[]): void {
    const [w1, b1, w2, b2, w3, b3, w4, b4, w5, b5] = weights
    this.conv1.setWeights([w1, b1])
    this.conv2.setWeights([w2, b2])
    this.conv3.setWeights([w3, b3])
    this.conv4.setWeights([w4, b4])
    this.conv5.setWeights([w5, b5])
  }

  call(x: any) {
    const concat = (list: any[]) =>
      tf.layers.concatenate({ axis: 3 }).apply(list)
    let x1 = this.lrelu(this.conv1!.apply(x))
    let x2 = this.lrelu(this.conv2!.apply(concat([x, x1])))
    let x3 = this.lrelu(this.conv3!.apply(concat([x, x1, x2])))
    let x4 = this.lrelu(this.conv4!.apply(concat([x, x1, x2, x3])))
    let x5 = this.conv5!.apply(concat([x, x1, x2, x3, x4]))
    // Scale the residual by 0.2 and add to the input
    if (x instanceof tf.Tensor) {
      tf.dispose([x1, x2, x3, x4])
      // @ts-ignore
      const out1 = x5.mul(0.2)
      const out2 = tf.add(x, out1)
      tf.dispose([x5, out1])
      return out2
    }
    return x
  }
  static get className() {
    return 'ResidualDenseBlock'
  }
}
tf.serialization.registerClass(ResidualDenseBlock)

// const model = new ResidualDenseBlock()
// const input = tf.randomNormal([1, 64, 64, 64]) // Example input tensor
// const output = model.apply(input) as tf.Tensor
// output.print()
